# sound_recognizer.py
import numpy as np
from scipy.io import wavfile
from scipy.fft import fft


def sound_recognize(wave, sample_rate=44100, target_frequency=[17000, 18000, 19000]):
    """
    识别声音中的目标频率
    """
    n = len(wave)
    fft_result = fft(wave)
    freqs = np.fft.fftfreq(n, 1 / sample_rate)

    half_n = n // 2
    fft_result = fft_result[:half_n]
    freqs = freqs[:half_n]

    posiblity_list = []
    for f in target_frequency:
        idx = np.argmin(np.abs(freqs - f))
        energy = np.sum(
            np.abs(fft_result[idx - 1 : idx + 2]) ** 2
        )  # 计算附近频率的能量
        posiblity_list.append(energy)

    result = [p / sum(posiblity_list) for p in posiblity_list]
    return result


def sound_recognize_from_file(file_name, target_frequency=[17000, 18000, 19000]):
    """
    从文件中读取声音并进行识别
    """
    sample_rate, wave = wavfile.read(file_name)
    if len(wave.shape) > 1:
        wave = wave[:, 0]  # 如果是立体声，取单声道
    return sound_recognize(wave, sample_rate, target_frequency)


if __name__ == "__main__":
    path = "sounds/sound.wav"
    target_frequency = [18500, 18700, 18900,19100]
    print(sound_recognize_from_file(path,target_frequency))
